# !/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time : 2023/4/24 10:12
# @Author : xyjiang1216
# @Email : jiangxuyan@nudt.edu.cn
# @File : layer_2_construct_stream.py
# @Version: OpenPlanner2.0
import importlib
import json
from random import choice, randint

import networkx as nx

from OpenPlanner.core.src.op_class import OPClass
from OpenPlanner.debug.debug_utils import DEBUG_PRINT
from OpenPlanner.util.json_utils import write_json
from OpenPlanner.util.topo_graph_utils import read_graph
from OpenPlanner.core.api.algo_input_class import StreamInfo



class StreamConstructionModule:
    def __init__(self,
                 op_obj: OPClass):
        self.op_obj = op_obj
        self.stream_hook_func = None

    def construct_stream_json(self):
        """
        :return:
        """
        if self.op_obj.stream_instance_flag == 'False':
            stream_filepath = \
                self.op_obj.stream_filepath
            try:
                fd = open(stream_filepath, 'r')
            except IOError:
                print('Error: 无需实例化流量，'
                      '但用户自定义的'
                      '流量文件%s'
                      '不存在或无法打开，'
                      '请检查文件路径' % stream_filepath)
                exit(0)
            else:
                fd.close()
        elif self.op_obj.stream_instance_flag == 'True':
            '''从拓扑文件中读出拓扑图'''
            topo_filepath = self.op_obj.topo_filepath
            graph = read_graph(topo_filepath)

            '''在图上随机生成流量的路由路径'''
            route_path_set = self._random_route_path(graph,
                                                     max_hop=7,
                                                     strategy='shortest_path')

            '''随机生成流量的参数'''
            stream_set = \
                self._random_stream_param(route_path_set)

            '''写流量文件'''
            write_json(stream_set,
                       self.op_obj.stream_filepath)
        else:
            print("Error: 是否需要实例化流量的标识错误，"
                  "只能是True或False，"
                  "请检查配置文件")
            exit(0)

    def _random_route_path(self, graph, max_hop=7,
                           strategy='shortest_path'):
        # 1. 找到所有终端节点
        # 对于双向图，只要节点的度为2，
        # 就可以认为它是一个终端节点
        end_set = []
        sw_set = []
        for item in graph.degree:
            if item[1] <= 2:
                end_set.append(item[0])
            else:
                sw_set.append(item[0])

        stream_num = self.op_obj.stream_param.stream_num
        route_path_set = []
        stream = 0
        while stream < stream_num:
            # for stream in range(stream_num):
            src_es_id = choice(end_set)
            dst_es_id = choice(end_set)
            while dst_es_id == src_es_id:
                dst_es_id = choice(end_set)

            shortest_path_in_node_id = nx.shortest_path(graph, src_es_id, dst_es_id)
            # 将路由路径长度限制为7
            if len(shortest_path_in_node_id) > max_hop + 1:
                continue
            # shortest_path_in_node_id是node_id的集合
            # 将shortest_path_in_node_id转换成link_id的集合
            shortest_path_in_link_id = []
            for src_node_id, dst_node_id in zip(shortest_path_in_node_id, shortest_path_in_node_id[1:]):
                link_id = graph[src_node_id][dst_node_id]['link_id']
                shortest_path_in_link_id.append(link_id)
            route_path_set.append(shortest_path_in_link_id)
            stream += 1

        return route_path_set

    def _random_stream_param(self, route_path_set):
        stream_num = self.op_obj.stream_param.stream_num
        if stream_num <= 0:
            print("Error: 配置文件中的流量数目（stream_num）"
                  "小于或等于0，"
                  "请更新配置文件")
            exit(0)
        size_set = self.op_obj.stream_param.size_set
        if not size_set:
            print("Error: 配置文件中的报文大小数组（size_set）"
                  "不能为空，"
                  "请更新配置文件")
            exit(0)
        period_set = self.op_obj.stream_param.period_set
        if not period_set:
            print("Error: 配置文件中的流量周期数组（period_set）"
                  "不能为空，"
                  "请更新配置文件")
            exit(0)
        user_defined_param = \
            self.op_obj.stream_param.user_defined_param
        stream_set = []
        for (stream_id, route) in \
                zip(range(stream_num), route_path_set):
            size = choice(size_set)
            index = randint(0, len(period_set) - 1)
            period = period_set[index]
            stream = {'stream_id': stream_id,
                      'size': size,
                      'period': period,
                      # 'latency_requirement': latency_requirement,
                      # 'jitter_requirement': jitter_requirement,
                      'route': route}
            # 判断是否有自定义参数，
            # 若有，添加自定义参数
            for item in user_defined_param:
                try:
                    stream.update({item: choice(user_defined_param[item])})
                except IndexError:
                    print("Error: 流量自定义参数"
                          "user_defined_param的"
                          "%s为空，请更新配置文件" % item)
            stream_set.append(stream)
        return stream_set

    # 钩子函数相关
    def _register_stream_hook(self, func):
        self.stream_hook_func = func

    def custom_stream_hook(self):
        base_module = self.op_obj.stream_hook.base_module_path
        func_name = self.op_obj.stream_hook.hook_func_name
        if base_module != '' and func_name != '':
            # 注册钩子
            # 判断模块是否存在
            try:
                module = importlib.import_module(base_module)
            except ModuleNotFoundError:
                print("ModuleNotFoundError:"
                      "流量钩子函数的%s模块不存在，"
                      "请检查钩子函数路径后更新配置文件" % base_module)
                exit(0)
            # 判断某个模块中是否存在某个函数
            try:
                func = getattr(module, func_name)
            except AttributeError:
                print("AttributeError:"
                      "模块%s内没有函数%s，"
                      "请检查拓扑钩子函数名称后更新配置文件"
                      % (base_module, func_name))
                exit(0)
            self._register_stream_hook(func)

        # 执行钩子函数
        if self.stream_hook_func is not None:
            kwargs = self.op_obj.stream_hook.hook_func_param
            self.stream_hook_func(self.op_obj, **kwargs)
        else:
            DEBUG_PRINT("No hook after topo construction")
            pass

